package com.star.spring.ratelimit.aspect;

import java.lang.reflect.Method;

import jakarta.servlet.http.HttpServletRequest;

import lombok.extern.slf4j.Slf4j;

import com.star.spring.ratelimit.annotation.RateLimit;
import com.star.spring.ratelimit.exception.RateLimitException;
import com.star.spring.ratelimit.limiter.LocalSlidingWindowLimiter;
import com.star.spring.ratelimit.limiter.LocalTokenBucketLimiter;
import com.star.spring.ratelimit.limiter.RedisFixedWindowLimiter;
import com.star.spring.ratelimit.limiter.RedisSlidingWindowLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.Order;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.Expression;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

/**
 * 限流切面
 *
 * @author star
 * @since 2025-01-24
 */
@Slf4j
@Aspect
@Component
@Order(2) // 第二个执行
public class RateLimitAspect {

    @Autowired
    private LocalTokenBucketLimiter localTokenBucketLimiter;

    @Autowired
    private LocalSlidingWindowLimiter localSlidingWindowLimiter;

    @Autowired
    private RedisFixedWindowLimiter redisFixedWindowLimiter;

    @Autowired
    private RedisSlidingWindowLimiter redisSlidingWindowLimiter;

    private final ExpressionParser parser = new SpelExpressionParser();

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        if (!rateLimit.enabled()) {
            return joinPoint.proceed();
        }

        // 生成限流key
        String key = generateKey(joinPoint, rateLimit);

        // 根据限流类型执行限流逻辑
        boolean acquired = false;
        switch (rateLimit.type()) {
            case LOCAL_TOKEN_BUCKET:
                acquired = localTokenBucketLimiter.tryAcquire(key, rateLimit.count(), rateLimit.timeWindow());
                break;
            case LOCAL_SLIDING_WINDOW:
                acquired = localSlidingWindowLimiter.tryAcquire(key, rateLimit.count(), rateLimit.timeWindow());
                break;
            case REDIS_FIXED_WINDOW:
                acquired = redisFixedWindowLimiter.tryAcquire(key, rateLimit.count(), rateLimit.timeWindow());
                break;
            case REDIS_SLIDING_WINDOW:
                acquired = redisSlidingWindowLimiter.tryAcquire(key, rateLimit.count(), rateLimit.timeWindow(), 1);
                break;
            default:
                log.warn("Unknown rate limit type: {}", rateLimit.type());
                return joinPoint.proceed();
        }

        if (!acquired) {
            log.warn("Rate limit exceeded - key: {}, type: {}", key, rateLimit.type());
            throw new RateLimitException(rateLimit.message());
        }

        return joinPoint.proceed();
    }

    /**
     * 生成限流key
     */
    private String generateKey(ProceedingJoinPoint joinPoint, RateLimit rateLimit) {
        String key = rateLimit.key();

        if (key.isEmpty()) {
            // 默认使用类名+方法名
            String className = joinPoint.getTarget().getClass().getSimpleName();
            String methodName = joinPoint.getSignature().getName();
            key = className + ":" + methodName;
        } else {
            // 解析SpEL表达式
            try {
                Expression expression = parser.parseExpression(key);
                EvaluationContext context = new StandardEvaluationContext();

                // 添加方法参数
                Method method = ((org.aspectj.lang.reflect.MethodSignature) joinPoint.getSignature()).getMethod();
                Object[] args = joinPoint.getArgs();
                String[] paramNames =
                        ((org.aspectj.lang.reflect.MethodSignature) joinPoint.getSignature()).getParameterNames();

                for (int i = 0; i < paramNames.length; i++) {
                    context.setVariable(paramNames[i], args[i]);
                }

                // 添加请求信息
                ServletRequestAttributes attributes =
                        (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
                if (attributes != null) {
                    HttpServletRequest request = attributes.getRequest();
                    context.setVariable("request", request);
                    context.setVariable("ip", getClientIp(request));
                    context.setVariable("userAgent", request.getHeader("User-Agent"));
                }

                Object result = expression.getValue(context);
                key = result != null ? result.toString() : key;
            } catch (Exception e) {
                log.warn("Failed to parse SpEL expression: {}, using original key", key, e);
            }
        }

        return key;
    }

    /**
     * 获取客户端IP
     */
    private String getClientIp(HttpServletRequest request) {
        String xForwardedFor = request.getHeader("X-Forwarded-For");
        if (xForwardedFor != null && !xForwardedFor.isEmpty() && !"unknown".equalsIgnoreCase(xForwardedFor)) {
            return xForwardedFor.split(",")[0].trim();
        }

        String xRealIp = request.getHeader("X-Real-IP");
        if (xRealIp != null && !xRealIp.isEmpty() && !"unknown".equalsIgnoreCase(xRealIp)) {
            return xRealIp;
        }

        return request.getRemoteAddr();
    }
}
